# -*- coding: utf-8 -*-
# @time: 2024/6/4 17:07
# @file: RSJ
# @author: tyshixi08

from get_data.origin_data import *

# 获取转债收益率
def get_bond_return(code='866005.RI', start_date = month_ls()[0].replace('-','') + '01', end_date = month_ls()[-1].replace('-','') + '01'):
    df = get_price_change_rate(code, start_date, end_date)
    df = df.T
    df = df.stack()
    df = df.reset_index().rename(columns={0:'return'})
    df['date'] = pd.to_datetime(df['date']).dt.strftime('%Y-%m-%d')
    return df.sort_values('date')

def RSJ(code='866005.RI', start_date = month_ls()[0].replace('-','') + '01', end_date = month_ls()[-1].replace('-','') + '01', window = 10):
    df = get_bond_return(code, start_date, end_date)

    # 确保日期列是 datetime 类型
    df['date'] = pd.to_datetime(df['date'])

    # 设定 'date' 列为索引
    df.set_index('date', inplace=True)

    df = df.sort_values(by='date')

    grouped = df['return'].rolling(window=window)

    # 定义一个函数来计算大于0的收益率的平方和
    def positive_squared_sums(x):

        positive_data = np.array([r for r in x if r > 0])
        positive_squared_sum = np.sum(pow(positive_data, 2))

        return positive_squared_sum

    # 计算正向波动率
    positive_results = grouped.apply(positive_squared_sums)
    positive_results = positive_results.reset_index().rename(columns={'return': 'positive_RV'})

    # 定义一个函数来计算小于0的收益率的平方和
    def negative_squared_sums(x):

        negative_data = np.array([r for r in x if r <= 0])
        negative_squared_sum = np.sum(pow(negative_data, 2))

        return negative_squared_sum

    # 计算负向波动率
    negative_results = grouped.apply(negative_squared_sums)
    negative_results = negative_results.reset_index().rename(columns={'return': 'negative_RV'})

    # 定义一个函数来计算收益率的总平方和
    def all_squared_sums(x):

        all_squared_sum = np.sum(pow(x, 2))

        return all_squared_sum

    # 计算总波动率
    all_results = grouped.apply(all_squared_sums)
    all_results = all_results.reset_index().rename(columns={'return': 'all_RV'})

    df_RV = pd.merge(positive_results, negative_results, how='outer', on='date')
    df_RV = pd.merge(all_results, df_RV, how='outer', on='date')

    df = pd.merge(df, df_RV, how='outer', on='date').dropna()
    df['RSJ'] = (df['positive_RV'] - df['negative_RV']) / df['all_RV']

    return df.reset_index(drop=True)

def RSJ_median(code='866005.RI', start_date = month_ls()[0].replace('-','') + '01', end_date = month_ls()[-1].replace('-','') + '01', window = 10):

    df = RSJ(code, start_date, end_date, window)

    # 获取RSJ中位数
    df_RSJ = df.groupby('date')['RSJ'].median()
    df_RSJ = df_RSJ.reset_index().rename(columns={'RSJ': f'RSJ_{window}D'})

    return df_RSJ

def RSJ_data():
    df = pd.merge(RSJ_median(window = 5), RSJ_median(window = 10), how = 'outer', on = 'date')
    return df

# 数据存档
def RSJ_5D_save():
    excel_file_path = 'RSJ_5D.csv'
    #if os.path.exists('原始数据/' + excel_file_path):
    #    df = pd.read_csv('原始数据/' + excel_file_path)
    #    return df
    #else:
    save_CSV(RSJ_data()[['date', 'RSJ_5D']].sort_values('date').dropna(), 'get_data/原始数据/' + excel_file_path.split('.')[0])
    df = pd.read_csv('get_data/原始数据/' + excel_file_path)
    return df

# 数据存档
def RSJ_10D_save():
    excel_file_path = 'RSJ_10D.csv'
    #if os.path.exists('原始数据/' + excel_file_path):
    #    df = pd.read_csv('原始数据/' + excel_file_path)
    #    return df
    #else:
    save_CSV(RSJ_data()[['date', 'RSJ_10D']].sort_values('date').dropna(), 'get_data/原始数据/' + excel_file_path.split('.')[0])
    df = pd.read_csv('get_data/原始数据/' + excel_file_path)
    return df
